"""
3.	分水岭算法主要用于图像的分割，如果目标物体是连接在一起的，则分割起来会很困难，此时经常采用分水岭分割算法，会得到比较好的效果。
分水岭分割算法把图像看成一幅地形图，其中，亮度比较强的区域像素值较大，亮度暗的区域像素值比较小，通过寻找汇水盆地和分水岭界线对图像进行分割。
利用OpenCV+Python环境，代码按下列要求，完成分水岭算法的“硬币”图像分割。图像处理过程输出的5个窗口，参考下图所示（30分）
"""

import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt

spr = 2
spc = 3
spn = 0
plt.figure(figsize=(12, 6))


def my_imshow(title, img, color='gray'):
    """Utility function to show images."""
    global spn
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title(title)
    plt.axis('off')
    if 'gray' == color:
        plt.imshow(img, cmap='gray')
    elif 'bgr' == color:
        img = cv.cvtColor(img, cv.COLOR_BGR2RGB)
        plt.imshow(img)
    else:
        plt.imshow(img)


def my_debug(arr, name):
    """Utility function for debugging."""
    print(f'{name}: {arr.dtype}, {arr.shape}, min = {arr.min()}, max = {arr.max()}, mean = {arr.mean()}')


# ①	读取图像文件（water_coins.jpg）记为src，备份图像文件记为img，并显示原始图像
path = 'data/images/water_coins.jpg'
src = cv.imread(path, cv.IMREAD_COLOR)
img = src.copy()
my_imshow('original', img, 'bgr')

# ②	利用形态学相关技术，消除图像中的噪点和空洞，标记确定的背景sure_bg，并显示
gray = cv.cvtColor(src, cv.COLOR_BGR2GRAY)
ret, bin = cv.threshold(gray, 0, 255, cv.THRESH_OTSU + cv.THRESH_BINARY_INV)
kernel = np.ones((3, 3), dtype=np.uint8)
opening = cv.morphologyEx(bin, cv.MORPH_OPEN, kernel, iterations=2)
sure_bg = cv.morphologyEx(opening, cv.MORPH_CLOSE, kernel, iterations=1)
my_imshow('sure_bg', sure_bg)

# ③	利用OpenCV库函数，进行距离变换，并显示dist_transform
dist_transform = cv.distanceTransform(sure_bg, cv.DIST_L2, 5)
cv.normalize(dist_transform, dist_transform, 0, 255, cv.NORM_MINMAX)
dist_transform = dist_transform.astype(np.uint8)
my_imshow('dist_transform', dist_transform)

# ④	根据距离变换，设定适当的阈值，确定前景sure_fg，并显示
ret, sure_fg = cv.threshold(dist_transform, 255 * 0.9, 255, cv.THRESH_BINARY)
my_imshow('sure_fg', sure_fg)

# ⑤	实施分水岭算法，标签图像将会被修改，边界区域的标记将变为 -1，显示图像分割结果图
unknown = sure_bg - sure_fg
ret, marker = cv.connectedComponents(sure_fg)
marker += 1
marker[unknown == 255] = 0
cv.watershed(src, marker)
src[marker == -1] = (0, 255, 0)
my_imshow('splited', src, 'bgr')

# show all plotting
plt.show()
